#############################################################################################
#############################################################################################
# 
#   Open Source License/Disclaimer, Forecast Systems Laboratory NOAA/OAR/GSD, 
#   325 Broadway Boulder, CO 80305
#
#   This software is distributed under the Open Source Definition, which may be 
#   found at http://www.opensource.org/.
#   In particular, redistribution and use in source and binary forms, with or 
#   without modification, are permitted provided that the following conditions are met:
#
#   - Redistributions of source code must retain this notice, this list of 
#     conditions and the following disclaimer.
#
#   - Redistributions in binary form must provide access to this notice, this 
#     list of conditions and the following disclaimer, and the underlying source code.
#
#   - All modifications to this software must be clearly documented, and are 
#     solely the responsibility of the agent making the modifications.
#
#   - If significant modifications or enhancements are made to this software, 
#     the GSD Software Policy Manager (softwaremgr.fsl@noaa.gov) should be notified.
#
#   THIS SOFTWARE AND ITS DOCUMENTATION ARE IN THE PUBLIC DOMAIN AND ARE 
#   FURNISHED "AS IS." THE AUTHORS, THE UNITED STATES GOVERNMENT, ITS INSTRUMENTALITIES,
#   OFFICERS, EMPLOYEES, AND AGENTS MAKE NO WARRANTY, EXPRESS OR IMPLIED, AS TO
#   THE USEFULNESS OF THE SOFTWARE AND DOCUMENTATION FOR ANY PURPOSE. THEY ASSUME 
#   NO RESPONSIBILITY (1) FOR THE USE OF THE SOFTWARE AND DOCUMENTATION; OR (2) TO PROVIDE
#   TECHNICAL SUPPORT TO USERS. 
#############################################################################################
#############################################################################################


import inspect
import numpy as np
import numpy.random as rn
import os

# Import Data Assimilation class
from .data_assimilation_class import DataAssimilationClass

class DA_var(DataAssimilationClass):
    # This implements a data assimilation scheme in which only the
    # most recent data point is used to adjust the ensemble and
    # parametrization. The ensemble is then only propagated from one
    # data point to the next at each assimilation step. Propagation to
    # the horizon then occurs at the final step. Also, the data array
    # passed to the analysis step is perturbed by samples from normal
    # distribution controlled by the data covariance matrix. For
    # deterministic analysis schemes and some particle filters this
    # will not be necessary.

    # NOTE: This just takes care of how data is used in the
    # assimilation. For a specific data assimilation problem the
    # model2data map and data covariance matrix must be specified.

    # The class is initialized by reading in the data to arrays.
    # Data = (Ndata_pts)x(measurement size) numpy array
    # DataTime = (Ndata_pts) numpy vector
    # data_noise = either scalar standard deviation of data error 
    #              or vector of standard deviations in data error 
    #              at each data point
    # EnSize = Integer specifying size of ensemble
    # SimDim = Dimension of the simulation at each timestep
    # Horizon = Time to propagate ensemble forecast until after last data 
    #           point is assimilated
    # EnsembleClass = object of type EnsembleGeneratorClass(object) that will
    #                 be called to generate the ensemble
    # AnalysisClass = object of type AnalysisGeneratorClass(object) that will 
    #                 be called to form the analysis ensemble
    def __init__(self,domain,DataFilePath,EnSize,AnalysisClass):
        self.domain = domain
        self._analysis = AnalysisClass
        
        # Important sizes and dimensions
        self.EnSize = EnSize
        
    def obs_read(self,DataFilePath):
        # Define the list of observations
        self.obs_list = []

        # Define the observation error covariance
        self.Rinv = []

        # Define the list to record the observation availabe time step
        self.innov_time = []

        for time in self.domain.history:   
            if os.path.exists(os.path.join(DataFilePath,'OBS_'+self.domain.Name+'_'+str(int(time))+'.npz')) :
                print 'Reading observation file at %s' % str(time)
                npzfile = np.load(os.path.join(DataFilePath,'OBS_'+self.domain.Name+'_'+str(int(time))+'.npz'))
                self.obs_list.append(npzfile['OBS'])
                obs_error = npzfile['OBS_ERROR'] 
        
                # Construct Observational Error Rinv
                self.Rinv.append(self.domain.error_matrix ( obs_error ))

                # Record the time when the observation is available
                self.innov_time.append(time)

                    
    def DArun(self):
        # Data Assimilation process consists of 2 steps: 
        # 1.) Analysis generation 
        # 2.) return analysis
        
        # 1.) Analysis generation
        Analysis = self._analysis.create_analysis(self.domain, self.fg, self.Binv, self.obs_list, \
                                                  self.Rinv, self.innov_time)
        # 2.) return analysis
        return Analysis
             

